mcp/streamable: use event store to fix unbounded memory issues #335

samthanawalla · 2025-08-20T15:44:34Z

This CL utilizes the event store to write outgoing messages and removes the unbounded outgoing data structure.

This CL utilizes the event store to write outgoing messages and removes the unbounded outgoing data structure. For modelcontextprotocol#190

jba · 2025-08-20T16:41:44Z

mcp/streamable.go

+			for data, err := range c.eventStore.After(ctx, c.SessionID(), stream.id, lastIndex) {
+				if err != nil {
+					// Wait for session initialization before yielding.
+					if errors.Is(err, ErrUnknownSession) || errors.Is(err, ErrUnknownStream) {


Instead of doing it this way, I would avoid calling After at all if there is no session or stream.

If there is a session and stream and After returns one of these errors, I think it is a real error and should be yielded.

I am not sure there is an easy way to do that because the session and stream may exist but it may not exist in the event store yet.

Append could happen before or after the After call which is why we need After to report the error to us.

You should only call After when the client sends Last-Event-ID. If they send it too early, the server should return an error. I don't understand the state where After is called before Append.

You should only call After when the client sends Last-Event-ID.

We can call After with an index of -1 to start writing from the beginning of the stream which allows us to simplify the logic even if last-event-id is not sent.

I don't understand the state where After is called before Append.

After is called when respondSSE is called which is disjoint from when Append is called in Write. These events can happen in any order which is why we case on ErrUnknownSession and ErrUnknownStream to skip to the logic below which waits for a stream signal.

Why don't we just call Append with nil data when we create the stream?

Here's the problem: I'd like the eventstore to be able to completely clean up the stream or session at will, and so when we get an unknown session or stream, we should fail this connection because it will never be recoverable.

findleyr · 2025-08-20T20:15:32Z

mcp/event.go

@@ -283,6 +283,12 @@ func (s *MemoryEventStore) Append(_ context.Context, sessionID string, streamID
 // index is no longer available.
 var ErrEventsPurged = errors.New("data purged")

+// ErrUnknownSession is the error that [EventStore.After] should return if the session ID is unknown.


Nit: s/return/wrap: we don't return this error value exactly, but one that wraps it.

findleyr · 2025-08-22T14:47:01Z

mcp/streamable.go


 			// If all requests have been handled and replied to, we should terminate this connection.
 			// "After the JSON-RPC response has been sent, the server SHOULD close the SSE stream."
 			// §6.4, https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#sending-messages-to-the-server
 			// We only want to terminate POSTs, and GETs that are replaying. The general-purpose GET
 			// (stream ID 0) will never have requests, and should remain open indefinitely.
-			if nOutstanding == 0 && !persistent {
+			if nOutstanding == 0 && !persistent && lastIndex >= int(stream.lastWriteIndex.Load()) {


I originally thought that this commit reintroduced the bug that was fixed in findleyr@f4a9396. However, I see that it probably doesn't, because of this atomic check.

I think it would be simpler to just move the check for nOutstanding above the After loop above. Then you don't need lastWriteIndex. WDYT?

I prefer to avoid atomics when there's already a synchronization mechanism (mu), because it's hard to reason about the relationship between the atomics and critical sections.

findleyr · 2025-08-22T14:57:39Z

mcp/streamable.go

+			for data, err := range c.eventStore.After(ctx, c.SessionID(), stream.id, lastIndex) {
+				if err != nil {
+					// Wait for session initialization before yielding.
+					if errors.Is(err, ErrUnknownSession) || errors.Is(err, ErrUnknownStream) {


Why don't we just call Append with nil data when we create the stream?

Here's the problem: I'd like the eventstore to be able to completely clean up the stream or session at will, and so when we get an unknown session or stream, we should fail this connection because it will never be recoverable.

jba

Sorry that this PR is confusing me so much.
I think I need to see the algorithm as a whole, without diffs.
So if you think it is correct after addressing my comments that it is correct, I'll approve and then read it separately.

jba · 2025-08-25T19:07:43Z

mcp/event.go

@@ -274,6 +274,12 @@ func (s *MemoryEventStore) Append(_ context.Context, sessionID string, streamID
 	// Purge before adding, so at least the current data item will be present.
 	// (That could result in nBytes > maxBytes, but we'll live with that.)
 	s.purge()
+
+	// An empty data slice signals that a stream has been registered.


"...signals that the stream has just been created."

This behavior should be documented in EventStore.After.

Actually, I'd rather have a new method for this, instead of overloading append. Start(ctx, sessionID, streamID)?

IIUC, the purpose of this is so that After will not return an error if it is called before the first Append, is that right?

jba · 2025-08-25T19:17:12Z

mcp/streamable.go

@@ -675,20 +653,27 @@ func (c *streamableServerConn) respondSSE(stream *stream, w http.ResponseWriter,
 // If the stream did not terminate normally, it is either because ctx was
 // cancelled, or the connection is closed: check the ctx.Err() to differentiate
 // these cases.
-func (c *streamableServerConn) messages(ctx context.Context, stream *stream, persistent bool) iter.Seq2[json.RawMessage, bool] {
-	return func(yield func(json.RawMessage, bool) bool) {
+func (c *streamableServerConn) messages(ctx context.Context, stream *stream, persistent bool, lastIndex int) iter.Seq2[json.RawMessage, error] {


doc the persistent and lastIndex args.

jba · 2025-08-25T19:17:26Z

mcp/streamable.go

 					return
 				}
+				// The stream exists, but does not contain any messages on the stream.


s/on the stream//

jba · 2025-08-25T19:18:32Z

mcp/streamable.go

 					return
 				}
+				// The stream exists, but does not contain any messages on the stream.
+				// Do not yield this data.
+				if data == nil {


I'm confused. If there are no messages, how did we get inside this loop? Shouldn't After yield nothing at all?

samthanawalla · 2025-08-25T19:23:30Z

Sorry that this PR is confusing me so much. I think I need to see the algorithm as a whole, without diffs. So if you think it is correct after addressing my comments that it is correct, I'll approve and then read it separately.

Sorry I prematurely sent a new commit- not ready for review yet! (The flow is very confusing, I agree)

mcp/streamable: use event store to fix unbounded memory issues

0f55f3e

This CL utilizes the event store to write outgoing messages and removes the unbounded outgoing data structure. For modelcontextprotocol#190

samthanawalla marked this pull request as ready for review August 20, 2025 15:47

samthanawalla requested review from findleyr and jba August 20, 2025 15:47

jba requested changes Aug 20, 2025

View reviewed changes

findleyr reviewed Aug 22, 2025

View reviewed changes

remove lastWriteIndex and register stream with event store when created

deda34e

jba requested changes Aug 25, 2025

View reviewed changes

samthanawalla marked this pull request as draft August 25, 2025 19:23

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

mcp/streamable: use event store to fix unbounded memory issues #335

mcp/streamable: use event store to fix unbounded memory issues #335

Uh oh!

samthanawalla commented Aug 20, 2025

Uh oh!

jba Aug 20, 2025

Uh oh!

samthanawalla Aug 20, 2025

Uh oh!

jba Aug 20, 2025

Uh oh!

samthanawalla Aug 21, 2025 •

edited

Loading

Uh oh!

findleyr Aug 22, 2025

Uh oh!

findleyr Aug 20, 2025

Uh oh!

findleyr Aug 22, 2025

Uh oh!

findleyr Aug 22, 2025

Uh oh!

jba left a comment

Uh oh!

jba Aug 25, 2025

Uh oh!

jba Aug 25, 2025

Uh oh!

jba Aug 25, 2025

Uh oh!

jba Aug 25, 2025

Uh oh!

jba Aug 25, 2025

Uh oh!

samthanawalla commented Aug 25, 2025

Uh oh!

Uh oh!

mcp/streamable: use event store to fix unbounded memory issues #335

Are you sure you want to change the base?

mcp/streamable: use event store to fix unbounded memory issues #335

Uh oh!

Conversation

samthanawalla commented Aug 20, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

samthanawalla Aug 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jba left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

samthanawalla commented Aug 25, 2025

Uh oh!

Uh oh!

samthanawalla Aug 21, 2025 •

edited

Loading